import argparse
from datetime import datetime
from pathlib import Path

import torch
import torchvision
from PIL import Image
from torch.utils.data import DataLoader

from models.encoder_decoder_conv_lstm import EncoderDecoderConvLSTM
from models.model import MyModel
from utils.data_load import SequenceDataset
from utils.device_utils import get_device
from utils.save_model.save_model_path import get_save_distributed_encoder_decoder_conv_lstm_model_path

device = get_device(target_gpu=0)

# 加载模型
def load_model(model_path):
    # 实例化模型结构
    model = EncoderDecoderConvLSTM().to(device)
    # 加载保存的模型权重
    model.load_state_dict(torch.load(model_path,weights_only=True))
    # 将模型设置为评估模式
    model.eval()
    return model

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="命令行工具")

    # 添加参数
    parser.add_argument("--csv_file", help="csv_file文件路径", default=r'/DATA/disk1/hu/weather/test/TestA.csv')
    parser.add_argument("--image_prefix_path", help="image 图片路径前缀",
                        default=r'/DATA/disk1/hu/weather/test/TestA/TestA/Radar')
    parser.add_argument("--batch_size", help="batch_size", default=80)
    parser.add_argument("--num_workers", help="num_workers", default=1)
    args = parser.parse_args()

    model = load_model(get_save_distributed_encoder_decoder_conv_lstm_model_path())
    # 定义数据变换（可选）
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
    ])
    # 创建Dataset和DataLoader
    dataset = SequenceDataset(
        csv_file=args.csv_file,
        image_prefix_path=args.image_prefix_path,
        transform=transform,
        sequence_length=40,
        split_ratio=0.5  # 前20张作为输入，后20张作为目标
    )

    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        prefetch_factor=10,
        persistent_workers=True,
        pin_memory=True  # 如果使用GPU，建议开启
    )

    with torch.no_grad():  # 禁用梯度计算，因为我们只进行前向传播
        x,y = (next(iter(dataloader)))
        output = model(x.to(device))
        print(output.shape)

        path_d=f'./image/predict/predict_image_{datetime.now().strftime("%Y-%m-%d_%H-%M")}'
        Path(path_d).mkdir(parents=True, exist_ok=True)

        # tensor_shape = output.shape
        # output = output.reshape(tensor_shape[0], tensor_shape[2], tensor_shape[1], tensor_shape[4], tensor_shape[3])
        # y = y.reshape(tensor_shape[0], tensor_shape[2], tensor_shape[1], tensor_shape[4], tensor_shape[3])

        output = output.permute(0,2,1,4,3)
        y = y.permute(0,2,1,4,3)

        for i in range(output.shape[0]):

            predict_data = ((output*255).cpu()[i][0][0]).byte().numpy()
            image = Image.fromarray(predict_data, mode='L')  # 'L' 模式表示灰度图
            image.save(f'{path_d}/predict_image_{i}.png')

            y_data = ((y*255).cpu()[i][0][0] ).byte().numpy()
            image = Image.fromarray(y_data, mode='L')  # 'L' 模式表示灰度图
            image.save(f'{path_d}/y_image_{i}.png')
    print('------------------------------')